package org.supermind.solr.search;

import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.queryParser.ParseException;
import org.apache.lucene.queryParser.QueryParser;
import org.apache.lucene.search.BooleanClause;
import org.apache.lucene.search.BooleanQuery;
import org.apache.lucene.search.Query;
import org.apache.solr.common.params.CommonParams;
import org.apache.solr.common.params.DefaultSolrParams;
import org.apache.solr.common.params.DisMaxParams;
import org.apache.solr.common.params.SolrParams;
import org.apache.solr.request.SolrQueryRequest;
import org.apache.solr.schema.IndexSchema;
import org.apache.solr.search.DisMaxQParser;
import org.apache.solr.search.QParser;
import org.apache.solr.search.QueryParsing;
import org.apache.solr.search.SolrQueryParser;
import org.apache.solr.util.SolrPluginUtils;

import java.util.ArrayList;
import java.util.Map;

/**
 * A watered-down version of dismax that only supports: bf, bq and qf.
 * Also preserves the use of SolrPluginUtils.partialEscape and SolrPluginUtils.stripIllegalOperators.
 */
public class BoostMaxQParser extends DisMaxQParser {
  protected Map<String, Float> queryFields;
  protected boolean escapeQuery;
  BoostMaxQueryParser lparser;

  public BoostMaxQParser(String qstr, SolrParams localParams, SolrParams params, SolrQueryRequest req) {
    super(qstr, localParams, params, req);
    this.escapeQuery = params.getBool("escapeQuery", true);
  }

  public Query parse() throws ParseException {
    SolrParams solrParams = localParams == null ? params : new DefaultSolrParams(localParams, params);
    queryFields = SolrPluginUtils.parseFieldBoosts(solrParams.getParams(DisMaxParams.QF));

    /* the main query we will execute.  we disable the coord because
     * this query is an artificial construct
     */
    BooleanQuery query = new BooleanQuery(true);

    addMainQuery(query, solrParams);
    addBoostQuery(query, solrParams);
    addBoostFunctions(query, solrParams);

    return query;
  }


  protected void addMainQuery(BooleanQuery query, SolrParams solrParams) throws ParseException {
//    Map<String, Float> phraseFields = SolrPluginUtils.parseFieldBoosts(solrParams.getParams(DisMaxParams.PF));
//    float tiebreaker = solrParams.getFloat(DisMaxParams.TIE, 0.0f);

    /* a parser for dealing with user input, which will convert
     * things to DisjunctionMaxQueries
     */
//    SolrPluginUtils.DisjunctionMaxQueryParser up = getParser(queryFields, DisMaxParams.QS, solrParams, tiebreaker);

    /* for parsing sloppy phrases using DisjunctionMaxQueries */
//    SolrPluginUtils.DisjunctionMaxQueryParser pp = getParser(phraseFields, DisMaxParams.PS, solrParams, tiebreaker);

    /* * * Main User Query * * */
    parsedUserQuery = null;
    String userQuery = getString();
    altUserQuery = null;

    if (userQuery == null || userQuery.trim().length() < 1) {
//       If no query is specified, we may have an alternate
      altUserQuery = getAlternateUserQuery(solrParams);
      query.add(altUserQuery, BooleanClause.Occur.MUST);
    } else {
      // There is a valid query string
      CharSequence cs = userQuery;
      if (escapeQuery) cs = SolrPluginUtils.partialEscape(cs);
      cs = SolrPluginUtils.stripUnbalancedQuotes(cs);
      cs = SolrPluginUtils.stripIllegalOperators(cs);
      userQuery = cs.toString();

      parsedUserQuery = getUserQuery(userQuery, solrParams);
      query.add(parsedUserQuery, BooleanClause.Occur.MUST);
    }
//      Query phrase = getPhraseQuery(userQuery, pp);
//      if (null != phrase) {
//        query.add(phrase, BooleanClause.Occur.SHOULD);
//      }
//    }
  }


  protected void addBoostQuery(BooleanQuery query, SolrParams solrParams) throws ParseException {
    boostParams = solrParams.getParams(DisMaxParams.BQ);
    //List<Query> boostQueries = SolrPluginUtils.parseQueryStrings(req, boostParams);
    boostQueries = null;
    if (boostParams != null && boostParams.length > 0) {
      boostQueries = new ArrayList<Query>();
      for (String qs : boostParams) {
        if (qs.trim().length() == 0) continue;
        Query q = subQuery(qs, null).parse();
        boostQueries.add(q);
      }
    }

    if (null != boostQueries) {
      for (Query f : boostQueries) {
        if (f instanceof BooleanQuery) {
          BooleanQuery b = (BooleanQuery) f;
          for (BooleanClause c : b.getClauses()) {
            if (c.isRequired()) c.setOccur(BooleanClause.Occur.SHOULD);
          }
          query.add(b, BooleanClause.Occur.SHOULD);
        } else {
          query.add(f, BooleanClause.Occur.SHOULD);
        }
      }
    }
  }

  protected Query getUserQuery(String userQuery, SolrParams solrParams) throws ParseException {
    initQueryParser();
    if (queryFields.size() > 0 && !userQuery.equals("*:*")) {
      final BooleanQuery bq = new BooleanQuery();
      for (Map.Entry<String, Float> entry : queryFields.entrySet()) {
        Query q = lparser.getFieldQuery(entry.getKey(), userQuery);
        if(q == null) return lparser.parse(userQuery); 
        if (entry.getValue() != null) q.setBoost(entry.getValue());
        bq.add(q, BooleanClause.Occur.SHOULD);
      }
      return bq;
    } else {
      return lparser.parse(userQuery);
    }
  }


  private void initQueryParser() {
    String defaultField = getParam(CommonParams.DF);
    if (defaultField == null) {
      defaultField = getReq().getSchema().getDefaultSearchFieldName();
    }
    lparser = new BoostMaxQueryParser(this, defaultField);

    // these could either be checked & set here, or in the SolrQueryParser constructor
    String opParam = getParam(QueryParsing.OP);
    if (opParam != null) {
      lparser.setDefaultOperator("AND".equals(opParam) ? QueryParser.Operator.AND : QueryParser.Operator.OR);
    } else {
      // try to get default operator from schema
      QueryParser.Operator operator = getReq().getSchema().getSolrQueryParser(null).getDefaultOperator();
      lparser.setDefaultOperator(null == operator ? QueryParser.Operator.OR : operator);
    }
  }

  public void setEscapeQuery(boolean escapeQuery) {
    this.escapeQuery = escapeQuery;
  }
}

class BoostMaxQueryParser extends SolrQueryParser {
  BoostMaxQueryParser(IndexSchema indexSchema, String s) {
    super(indexSchema, s);
  }

  BoostMaxQueryParser(QParser qParser, String s) {
    super(qParser, s);
  }

  BoostMaxQueryParser(QParser qParser, String s, Analyzer analyzer) {
    super(qParser, s, analyzer);
  }

  @Override public Query getFieldQuery(String field, String queryText) throws ParseException {
    Query q = super.getFieldQuery(field, queryText);
    return q;
  }
}